package com.aispeech.jarpc;

import com.aispeech.jarpc.message.ArpcMessage;
import com.aispeech.jarpc.message.EMsgType;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.IOException;
import java.io.InputStream;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;

import static com.aispeech.jarpc.message.ArpcMessage.MAX_MESSAGE_SIZE;
import static com.aispeech.jarpc.message.ArpcMessage.MESSAGE_HEADER_LENGTH;

@Slf4j
public class SyncSocket {
    private final static int MAX_RETRIES = 3;
    private Socket socket;

    public void connect(String host, int port, int timeout) throws IOException {
        String ip = "";
        UnknownHostException exception = null;

        for (int i = 0; i < MAX_RETRIES; i++) {
            try {
                ip = InetAddress.getAllByName(host)[0].getHostAddress();
                break;
            } catch (UnknownHostException unknownHostErr) {
                exception = unknownHostErr;
                ip = "";

                log.warn("[ SyncSocket.connect ] failed resolve host: {}, retry times: {}, error: {}", host, i, unknownHostErr.getMessage());
            }
        }

        if (StringUtils.isBlank(ip) && exception != null) throw exception;

//        log.info("begin connect to host: {}, {}:{}...", host, ip, port);
        socket = new Socket();

        socket.setTcpNoDelay(true);
        socket.setSoTimeout(timeout);

        socket.connect(new InetSocketAddress(ip, port), 3000);
    }

    public void close() {
        if (isConnected()) {
            try {
                socket.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        socket = null;
    }

    public boolean isConnected() {
        return (socket != null && socket.isConnected());
    }

    public void send(byte[] sessionId, ByteBuffer data) throws IOException {
        byte[] req = data.array();

        ByteBuffer byteBuffer = ByteBuffer.allocate(MESSAGE_HEADER_LENGTH + req.length + sessionId.length);
        byteBuffer.order(ByteOrder.LITTLE_ENDIAN);

        byteBuffer.put(EMsgType.PROCESS.getValue());
        byteBuffer.putInt(0);
        byteBuffer.putInt(req.length + sessionId.length);
        byteBuffer.put(req);
        byteBuffer.put(sessionId);

        byte[] payload = byteBuffer.array();

        socket.getOutputStream().write(payload);
        socket.getOutputStream().flush();
    }

    public ArpcMessage receive() throws ArpcException {
        try {
            InputStream inputStream = socket.getInputStream();

            byte[] buffer = new byte[MESSAGE_HEADER_LENGTH];
            int read = inputStream.read(buffer);

            if (read <= 0) throw new ArpcException(EArpcState.S_CONNECTION_BROKEN);

            ArpcMessage message = ArpcMessage.from(buffer);

            int bodyLength = message.getLength();
            if (bodyLength > 0) {
                if (bodyLength > MAX_MESSAGE_SIZE) throw new ArpcException(EArpcState.S_MSG_OVERLOAD);

                ByteBuffer body = ByteBuffer.allocate(bodyLength);
                buffer = new byte[8192];

                while (body.hasRemaining()) {
                    read = inputStream.read(buffer);
                    if (read <= 0) throw new ArpcException(EArpcState.S_CONNECTION_BROKEN);

                    body.put(buffer, 0, read);
                }

                message.setBody(body.array());
            }

            if (message.getMsgType() != EMsgType.STOP && message.getMsgType() != EMsgType.REDIRECT) {
                throw new ArpcException(EArpcState.S_INTERNAL_ERROR.getValue(), String.format("message: %s", message));
            }

            return message;
        } catch (IOException ex) {
            log.error("[ SyncSocket.receive ]", ex);
            throw new ArpcException(ex.getMessage());
        }
    }
}
